XLA_PYTHON_CLIENT_PREALLOCATE=false CUDA_VISIBLE_DEVICES=0 python klein_gordon3d.py --model=spinn --equation=klein_gordon3d --nc=64 --nc_test=100 --seed=111 --lr=0.001 --epochs=50000 --mlp=modified_mlp --n_layers=4 --features=64 --r=32 --out_dim=1 --pos_enc=0 --k=2 --log_iter=1000 --plot_iter=10000